{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "31ef0f85",
   "metadata": {},
   "source": [
    "This is part 1 of a tutorial series. We recommend following them in order, starting with [Part 0: Welcome to `musica`](0.%20Welcome%20to%20MUSICA.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91c25502",
   "metadata": {},
   "source": [
    "# `tuv-x`: Standard Configurations\n",
    "\n",
    "The [TUV-x photolysis calculator](https://github.com/ncar/tuv-x) is available in `musica` as `tuv-x`. `tuv-x` can be used to calculate photolysis rate constants for a given set of conditions. It can also be used to calculate exposure metrics (\"dose rates\" in TUV-x lingo) based on the calculated radiation field.\n",
    "\n",
    "`tuv-x` can be configured for any number of photolysis and dose rate calculations. We also provide two standard configurations in the MUSICA Python package:\n",
    "- v5.4: A `tuv-x` configuration that corresponds to the original [TUV 5.4 model](https://www2.acom.ucar.edu/modeling/tuv-download)\n",
    "- vTS1/TSMLT: A `tuv-x` configuration for use with the MOZART TS1 and TSMLT chemical mechanisms\n",
    "\n",
    "If you've been following the tutorial series from the beginning, you may remember that `micm` (the ODE solver for chemistry) can solve one or more \"grid cells\" (i.e., well-mixed air masses), but it makes no assumptions about the size, shape, or arrangement of those grid cells. Not so with `tuv-x`! To calculate photolysis rate constants and dose rates, `tuv-x` first needs to solve the radiation field in each grid cell, and to do that, it assumes rectilinear grid cells, stacked vertically to form \"columns\". The horizontal dimensions of the grid cells are arbitrary, but the vertical dimensions must be known by `tuv-x`.\n",
    "\n",
    "There are also many more input conditions needed by `tuv-x` than `micm`, so in this first `tuv-x` tutorial, we'll use a set of standard conditions available in `musica`. In the next tutorial, we'll see how to set conditions via the API.\n",
    "\n",
    "## 1. $\\text{O}_3$ Photolysis Comparison\n",
    "\n",
    "To get started, we'll use the two standard TUV-x configurations to calculate the rate constants for O3 photolysis. We include rate constants for both $\\text{O}_3$ photolysis branches (forming $\\text{O(1D)}$ or $\\text{O(3P)}$)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e189d851",
   "metadata": {},
   "outputs": [],
   "source": [
    "from musica.tuvx import v54, vTS1\n",
    "from musica.tuvx import TUVX\n",
    "import xarray as xr\n",
    "\n",
    "v54_calculator: TUVX = v54.get_tuvx_calculator()\n",
    "vTS1_calculator: TUVX = vTS1.get_tuvx_calculator()\n",
    "\n",
    "v54_results: xr.Dataset = v54_calculator.run(sza=0.0, earth_sun_distance=1.0)\n",
    "vTS1_results: xr.Dataset = vTS1_calculator.run(sza=0.0, earth_sun_distance=1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "036ba393",
   "metadata": {},
   "source": [
    "Here, we first load the required dependencies. TUV-x returns results as XArray Datasets, so we also import XArray.\n",
    "\n",
    "The `get_tuvx_calculator()` functions return fully configured TUV-x instances, loaded with specific rate constant and dose rate algorithms, along with a set of standard environmental conditions (we'll look at these more in a later section). The v5.4 and vTS1/TSMLT configurations include the same set of standard environmental conditions to facilitate comparisons. The environmental conditions and results from TUV-x are for a vertical atmosphere column. The standard conditions use a column with 120 layers of 1-km height per layer. TUV-x can be configured to work with columns with any number of layers.\n",
    "\n",
    "The `TUVX::run()` function calculates photolysis rate constants and dose rates by first solving the radiation field for the column, and then applying the configured photolysis and dose rate algorithms to the calculated radiation field. The two arguments to the `TUVX::run()` function set the solar zenith angle (`sza`) \\[radians\\] and the Earth–Sun Distance(`earth_sun_distance`) \\[AU\\].\n",
    "\n",
    "Now, let's compare the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf25af75",
   "metadata": {},
   "outputs": [],
   "source": [
    "# the labels for the O3 reactions are different in the two configurations\n",
    "jo3_o1d_v54_label = \"O3+hv->O2+O(1D)\"\n",
    "jo3_o3p_v54_label = \"O3+hv->O2+O(3P)\"\n",
    "jo3_o1d_vTS1_label = \"jo3_a\"\n",
    "jo3_o3p_vTS1_label = \"jo3_b\"\n",
    "\n",
    "# get the height and O3 profiles for both sets of results\n",
    "height_v54 = v54_results[\"vertical_edge\"].values\n",
    "jo3_o1d_v54 = v54_results[\"photolysis_rate_constants\"].sel(reaction=jo3_o1d_v54_label)\n",
    "jo3_o3p_v54 = v54_results[\"photolysis_rate_constants\"].sel(reaction=jo3_o3p_v54_label)\n",
    "height_vTS1 = vTS1_results[\"vertical_edge\"].values\n",
    "jo3_o1d_vTS1 = vTS1_results[\"photolysis_rate_constants\"].sel(reaction=jo3_o1d_vTS1_label)\n",
    "jo3_o3p_vTS1 = vTS1_results[\"photolysis_rate_constants\"].sel(reaction=jo3_o3p_vTS1_label)\n",
    "\n",
    "# plot the comparison\n",
    "import matplotlib.pyplot as plt\n",
    "plt.figure(figsize=(8, 6))\n",
    "plt.plot(jo3_o1d_v54, height_v54, label=\"TUV-x v5.4 O3->O(1D)\")\n",
    "plt.plot(jo3_o3p_v54, height_v54, label=\"TUV-x v5.4 O3->O(3P)\")\n",
    "plt.plot(jo3_o1d_vTS1, height_vTS1, label=\"TUV-x TS1/TSMLT O3->O(1D)\")\n",
    "plt.plot(jo3_o3p_vTS1, height_vTS1, label=\"TUV-x TS1/TSMLT O3->O(3P)\")\n",
    "plt.xlabel(\"Photolysis Rate Constant (s$^{-1}$)\")\n",
    "plt.ylabel(\"Height (km)\")\n",
    "plt.title(\"O3 Photolysis Rate Constants from TUV-x Configurations\")\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.show()\n",
    "\n",
    "# plot the differences\n",
    "plt.figure(figsize=(8, 6))\n",
    "plt.plot(jo3_o1d_v54 - jo3_o1d_vTS1, height_v54, label=\"Difference in O3->O(1D)\")\n",
    "plt.plot(jo3_o3p_v54 - jo3_o3p_vTS1, height_v54, label=\"Difference in O3->O(3P)\")\n",
    "plt.xlabel(\"Photolysis Rate Constant Difference (s$^{-1}$)\")\n",
    "plt.ylabel(\"Height (km)\")\n",
    "plt.title(\"Differences in O3 Photolysis Rate Constants between TUV-x Configurations\")\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b38f04e3",
   "metadata": {},
   "source": [
    "The differences between results for the two configurations are relatively small, and in fact both the v5.4 and vTS1/TSMLT configurations use the same configuration for $\\text{O}_3$ photolysis:\n",
    "\n",
    "```json\n",
    "         {\n",
    "            \"name\": \"O3+hv->O2+O(1D)\",\n",
    "            \"cross section\": {\n",
    "               \"netcdf files\": [\n",
    "                 { \"file path\": \"data/cross_sections/O3_1.nc\" },\n",
    "                 { \"file path\": \"data/cross_sections/O3_2.nc\" },\n",
    "                 { \"file path\": \"data/cross_sections/O3_3.nc\" },\n",
    "                 { \"file path\": \"data/cross_sections/O3_4.nc\" }\n",
    "               ],\n",
    "               \"type\": \"O3\",\n",
    "               \"__output\": true\n",
    "            },\n",
    "            \"quantum yield\": {\n",
    "               \"type\": \"O3+hv->O2+O(1D)\"\n",
    "            }\n",
    "         },\n",
    "         {\n",
    "            \"name\": \"O3+hv->O2+O(3P)\",\n",
    "            \"cross section\": {\n",
    "               \"netcdf files\": [\n",
    "                 { \"file path\": \"data/cross_sections/O3_1.nc\" },\n",
    "                 { \"file path\": \"data/cross_sections/O3_2.nc\" },\n",
    "                 { \"file path\": \"data/cross_sections/O3_3.nc\" },\n",
    "                 { \"file path\": \"data/cross_sections/O3_4.nc\" }\n",
    "               ],\n",
    "               \"type\": \"O3\",\n",
    "               \"__output\": true\n",
    "            },\n",
    "            \"quantum yield\": {\n",
    "               \"type\": \"O3+hv->O2+O(3P)\",\n",
    "               \"__output\": true\n",
    "            }\n",
    "         }\n",
    "```\n",
    "\n",
    "The small differences we do see can be attributed to the differences in the 156- and 102-bin wavelength grids of the v5.4 and vTS1/TSMLT configurations, respectively. Solving the radiation field on differently shaped wavelength grids is expected to lead to differences in the results.\n",
    "\n",
    "Now, let's look at results for a photolysis rate constant that is configured differently in v5.4 and vTS1/TSMLT, $\\text{N}_2\\text{O}_5$ photolysis:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b4f4aa9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# the labels for the N2O5 reactions are different in the two configurations\n",
    "jn2o5_v54_label = \"N2O5+hv->NO2+NO3\"\n",
    "jn2o5a_vTS1_label = \"jn2o5_a\" # N2O5 + hv -> NO2 + NO3\n",
    "jn2o5b_vTS1_label = \"jn2o5_b\" # N2O5 + hv -> NO + O + NO3\n",
    "\n",
    "# get the height and O3 profiles for both sets of results\n",
    "height_v54 = v54_results[\"vertical_edge\"].values\n",
    "jn2o5_v54 = v54_results[\"photolysis_rate_constants\"].sel(reaction=jn2o5_v54_label)\n",
    "height_vTS1 = vTS1_results[\"vertical_edge\"].values\n",
    "jn2o5a_vTS1 = vTS1_results[\"photolysis_rate_constants\"].sel(reaction=jn2o5a_vTS1_label)\n",
    "jn2o5b_vTS1 = vTS1_results[\"photolysis_rate_constants\"].sel(reaction=jn2o5b_vTS1_label)\n",
    "\n",
    "# plot the comparison\n",
    "import matplotlib.pyplot as plt\n",
    "plt.figure(figsize=(8, 6))\n",
    "plt.plot(jn2o5_v54, height_v54, label=\"TUV-x v5.4 N2O5->NO2+NO3\")\n",
    "plt.plot(jn2o5a_vTS1, height_vTS1, label=\"TUV-x TS1/TSMLT N2O5->NO2+NO3\")\n",
    "plt.plot(jn2o5b_vTS1, height_vTS1, label=\"TUV-x TS1/TSMLT N2O5->NO+O+NO3\")\n",
    "plt.xlabel(\"Photolysis Rate Constant (s$^{-1}$)\")\n",
    "plt.ylabel(\"Height (km)\")\n",
    "plt.title(\"N2O5 Photolysis Rate Constants from TUV-x Configurations\")\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.show()\n",
    "\n",
    "# plot the difference between the single v5.4 rate constant and the sum of the two TS1/TSMLT rate constants\n",
    "plt.figure(figsize=(8, 6))\n",
    "plt.plot(jn2o5_v54 - (jn2o5a_vTS1 + jn2o5b_vTS1), height_v54, label=\"Difference in N2O5 Photolysis Rate Constants\")\n",
    "plt.xlabel(\"Photolysis Rate Constant Difference (s$^{-1}$)\")\n",
    "plt.ylabel(\"Height (km)\")\n",
    "plt.title(\"Differences in net N2O5 Photolysis Rate Constants between TUV-x Configurations\")\n",
    "plt.legend()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2be4c26",
   "metadata": {},
   "source": [
    "Note that the v5.4 configuration only includes a single pathway for $\\text{N}_2\\text{O}_5$ photolysis, whereas the vTS1/TSMLT includes two. The rate constant for the single pathway in v5.4 is approximately equal to the combination of the two pathways in vTS1/TSMLT."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d54f26d6",
   "metadata": {},
   "source": [
    "## Dose Rates\n",
    "\n",
    "Let's plot a couple dose rate outputs, as well. The vTS1/TSMLT configuration doesn't include dose rates, as it's primarily designed for use with the MOZART TS1/TSMLT chemical mechanisms, so we'll just plot rates from the v5.4 configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45dd0dc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dose rate labels\n",
    "uva_label = \"UV-A, 315-400 nm\"\n",
    "uvb_label = \"UV-B, 280-315 nm\"\n",
    "\n",
    "# dose rates from the v5.4 configuration\n",
    "uva_v54 = v54_results[\"dose_rates\"].sel(dose_rate=uva_label)\n",
    "uvb_v54 = v54_results[\"dose_rates\"].sel(dose_rate=uvb_label)\n",
    "\n",
    "# plot the dose rate profiles\n",
    "plt.figure(figsize=(8, 6))\n",
    "plt.plot(uva_v54, height_v54, label=\"TUV-x v5.4 UV-A (315-400 nm) Dose Rate\")\n",
    "plt.plot(uvb_v54, height_v54, label=\"TUV-x v5.4 UV-B (280-315 nm) Dose Rate\")\n",
    "plt.xlabel(\"Dose Rate (W/m$^2$-nm)\")\n",
    "plt.ylabel(\"Height (km)\")\n",
    "plt.title(\"UV Dose Rates from TUV-x v5.4 Configuration\")\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv (3.12.3)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}